{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "321adc7f-f386-4dbf-bb9c-d5f789421f78",
   "metadata": {},
   "source": [
    "# Build and Deploy MCP Toolbox for Databases\n",
    "\n",
    "> IMPORTANT: This is the second notebook in the lab. The notebooks  build on top of each other, so be sure to run the preceding notebooks, in order, before running this one. Start your journey building ADK Agents with MCP Toolbox [here](./1_setup_and_explore_databases.ipynb). \n",
    "\n",
    "## Overview\n",
    "\n",
    "MCP Toolbox for Databases is an open source MCP server that enables you to develop Database tools for agents easier, faster, and more securely by handling the complexities such as connection pooling, authentication, and more. Toolbox supports [16+ database engines](https://googleapis.github.io/genai-toolbox/resources/sources/), including the Spanner and AlloyDB instances we provisioned for this lab.\n",
    "\n",
    "Toolbox sits between your application’s orchestration framework and your database, providing a control plane that is used to modify, distribute, or invoke tools. It simplifies the management of your tools by providing you with a centralized location to store and update tools, allowing you to share tools between agents and applications and update those tools without necessarily redeploying your application.\n",
    "\n",
    "![MCP Toolbox Architecture](img/toolbox_architecture.png)\n",
    "\n",
    "In this notebook, you will learn how to:\n",
    "\n",
    "- Define and Configure Tools: Create a `tools.yaml` file to define data sources (AlloyDB and Spanner) and related tools using SQL queries.\n",
    "- Securely Deploy MCP Toolbox: Deploy the Toolbox to a secure Cloud Run environment within a private VPC, using Secret Manager to handle sensitive credentials and configurations.\n",
    "- Test and Update Tools: Manually test the deployed tools using the toolbox-core library and then update them to include parameters for more dynamic and secure queries. This highlights how you can iterate on and improve your tools without redeploying your entire application.\n",
    "\n",
    "The architecture provisioned by this notebook emphasizes enterprise-readiness for AI agents by demonstrating how to build, deploy, and manage database tools in a secure and scalable manner, leveraging MCP Toolbox to handle complexities like connection pooling and authentication.\n",
    "\n",
    "### Terraform Resources\n",
    "\n",
    "The following pre-requisite resources were created for you by Terraform. See the [main.tf](../terraform/main.tf) file for more details on the environment configuration:\n",
    "\n",
    "- Custom VPC (demo-vpc): The private network where the Cloud Run service for MCP Toolbox is deployed, ensuring it is not exposed to the public internet.\n",
    "- AlloyDB Cluster & Instance: The AlloyDB database that serves as a data source for one of the Toolbox tools.\n",
    "- Spanner Instance & Database: The Spanner database that serves as another data source for the Toolbox.\n",
    "- IAM Service Account (toolbox-service-account): A dedicated service account with the minimum necessary permissions for the MCP Toolbox service to access the required databases and secrets.\n",
    "\n",
    "\n",
    "### Google Cloud Services Used in the Notebook\n",
    "\n",
    "This notebook makes use of the following Google Cloud services:\n",
    "\n",
    "- MCP Toolbox for Databases (Toolbox): An open-source MCP (Model Context Protocol) server that allows developers to connect gen AI agents to enterprise data easily and securely.\n",
    "- Cloud Run: To deploy and host the MCP Toolbox application as a secure, scalable service.\n",
    "- Secret Manager: To securely store and manage the tools.yaml configuration file and the AlloyDB database password.\n",
    "- IAM (Identity and Access Management): To manage the permissions for the toolbox-service-account service account.\n",
    "- AlloyDB: Used as a data source for a tool that retrieves transaction information.\n",
    "- Spanner: Used as a data source for a tool that retrieves account transfer information.\n",
    "- Vertex AI Workbench: The environment where this notebook is executed.\n",
    "\n",
    "### Logical Flow of the Notebook\n",
    "\n",
    "The notebook is structured to guide you from initial configuration to a fully deployed and tested MCP Toolbox instance:\n",
    "- Basic Setup: Configures the environment by defining variables, connecting to the Google Cloud project, setting up logging, and installing the required Python libraries.\n",
    "- Define `tools.yaml`: You will create the core configuration file for MCP Toolbox. This involves defining the AlloyDB and Spanner data sources, creating simple tools to query them, and grouping these tools into a toolset.\n",
    "- Secure Storage: The `tools.yaml` file and the AlloyDB password are securely stored in Secret Manager. This is a critical security practice to avoid hardcoding sensitive information.\n",
    "- Deploy to Cloud Run: The MCP Toolbox container is deployed to Cloud Run with a strong security posture, including deployment into a private VPC, using a dedicated service account, and mounting secrets from Secret Manager.\n",
    "- Test Tools: The initial static tools are tested to ensure the deployment was successful and that the Toolbox can connect to the databases. This is a useful practice to include in integration tests and CI/CD pipelines.\n",
    "- Update and Redeploy: The tools are updated to be more dynamic by adding parameters to the SQL queries. The updated `tools.yaml` is then written back to Secret Manager, and the Cloud Run service is redeployed to reflect the changes. This demonstrates the iterative development workflow enabled by MCP Toolbox.\n",
    "- Test Parameterized Tools: The updated, parameterized tools are tested to confirm they work as expected, highlighting how to pass arguments to your tools."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0927ae02-125e-4d3e-a885-1e437e6bddb6",
   "metadata": {},
   "source": [
    "## Basic Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "842a3ca3-4f62-4190-bbb6-10e7da9b13c8",
   "metadata": {},
   "source": [
    "### Define Notebook Variables\n",
    "\n",
    "Update the `project_id` and `region` variables below to match your environment. You can use defaults for the rest of the project variables. \n",
    "\n",
    "You will be prompted for two passwords:\n",
    "1. The AlloyDB admin password for the `postgres` user. This is the password you defined when provisioning the Terraform environment.\n",
    "2. The AlloyDB password for the least-privilege `toolbox_user` account you created in the first notebook. This is the user that the MCP Toolbox service will use to connect to the database. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4fa431b-2efd-4db1-89e1-b97b4fc722df",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Project variables\n",
    "project_id = \"my-project\"\n",
    "region = \"my-region\"\n",
    "vpc = \"demo-vpc\"\n",
    "gcs_bucket_name = f\"project-files-{project_id}\"\n",
    "\n",
    "# AlloyDB variables\n",
    "alloydb_cluster = \"my-alloydb-cluster\"\n",
    "alloydb_instance = \"my-alloydb-instance\"\n",
    "alloydb_database = \"finance\"\n",
    "alloydb_password = input(\"Please enter the password for the AlloyDB 'postgres' database user: \")\n",
    "alloydb_toolbox_user_password = input(\"Please enter the password for the AlloyDB 'toolbox_user' database user: \")\n",
    "\n",
    "# Spanner variables\n",
    "spanner_instance = \"my-spanner-instance\"\n",
    "spanner_database = \"finance-graph\"\n",
    "session = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01db8a8b-cae4-4978-bb9a-e2e5a8465696",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set env variable to suppress annoying system warnings when running shell commands\n",
    "%env GRPC_ENABLE_FORK_SUPPORT=1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62bc39dc-f6ee-455e-853a-643c12918c2f",
   "metadata": {},
   "source": [
    "### Connect to your Google Cloud Project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d972baa-ffa7-47c7-a965-fdaa114262b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Configure gcloud.\n",
    "!gcloud config set project {project_id}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38a124be-854e-4465-bab2-44ea855d9256",
   "metadata": {},
   "source": [
    "### Configure Logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66056e25-5e3a-47e3-a842-8fe6fadd9a33",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "# Configure the root logger to output messages with INFO level or above\n",
    "logging.basicConfig(level=logging.INFO, stream=sys.stdout, format='%(asctime)s[%(levelname)5s][%(name)14s] - %(message)s',  datefmt='%H:%M:%S', force=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b72eb606-f634-424f-be79-c1b75b1b1e14",
   "metadata": {},
   "source": [
    "### Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7c2e387-9dfc-49c2-8a32-2c54d7eebcc2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! pip install --quiet google-cloud-storage==2.19.0 \\\n",
    "                      toolbox-core==0.2.1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db5cf49a-a04d-429c-a5f3-2df8e068fc81",
   "metadata": {},
   "source": [
    "### Define Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "349aa29e-ccbb-4aef-b51b-20d9130c1e53",
   "metadata": {},
   "source": [
    "#### REST API Helper Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a0ad2b-fed0-434f-baa7-69247ee476c7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "import google.auth\n",
    "import json\n",
    "\n",
    "# Get an access token based upon the current user\n",
    "creds, _ = google.auth.default()\n",
    "authed_session = google.auth.transport.requests.AuthorizedSession(creds)\n",
    "access_token=creds.token\n",
    "\n",
    "if project_id:\n",
    "  authed_session.headers.update({\"x-goog-user-project\": project_id}) # Required to workaround a project quota bug\n",
    "\n",
    "def rest_api_helper(\n",
    "    url: str,\n",
    "    http_verb: str,\n",
    "    request_body: dict = None,\n",
    "    params: dict = None,\n",
    "    session: requests.Session = authed_session,\n",
    "  ) -> dict:\n",
    "  \"\"\"Calls a REST API using a pre-authenticated requests Session.\"\"\"\n",
    "\n",
    "  headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "  try:\n",
    "\n",
    "    if http_verb == \"GET\":\n",
    "      response = session.get(url, headers=headers, params=params)\n",
    "    elif http_verb == \"POST\":\n",
    "      response = session.post(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"PUT\":\n",
    "      response = session.put(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"PATCH\":\n",
    "      response = session.patch(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"DELETE\":\n",
    "      response = session.delete(url, headers=headers, params=params)\n",
    "    else:\n",
    "      raise ValueError(f\"Unknown HTTP verb: {http_verb}\")\n",
    "\n",
    "    # Raise an exception for bad status codes (4xx or 5xx)\n",
    "    response.raise_for_status()\n",
    "\n",
    "    # Check if response has content before trying to parse JSON\n",
    "    if response.content:\n",
    "        return response.json()\n",
    "    else:\n",
    "        return {} # Return empty dict for empty responses (like 204 No Content)\n",
    "\n",
    "  except requests.exceptions.RequestException as e:\n",
    "      # Catch potential requests library errors (network, timeout, etc.)\n",
    "      # Log detailed error information\n",
    "      print(f\"Request failed: {e}\")\n",
    "      if e.response is not None:\n",
    "          print(f\"Request URL: {e.request.url}\")\n",
    "          print(f\"Request Headers: {e.request.headers}\")\n",
    "          print(f\"Request Body: {e.request.body}\")\n",
    "          print(f\"Response Status: {e.response.status_code}\")\n",
    "          print(f\"Response Text: {e.response.text}\")\n",
    "          # Re-raise a more specific error or a custom one\n",
    "          raise RuntimeError(f\"API call failed with status {e.response.status_code}: {e.response.text}\") from e\n",
    "      else:\n",
    "          raise RuntimeError(f\"API call failed: {e}\") from e\n",
    "  except json.JSONDecodeError as e:\n",
    "      print(f\"Failed to decode JSON response: {e}\")\n",
    "      print(f\"Response Text: {response.text}\")\n",
    "      raise RuntimeError(f\"Invalid JSON received from API: {response.text}\") from e\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a97dc26-8e22-4505-9f9e-f702f4e26a88",
   "metadata": {},
   "source": [
    "## Setup MCP Toolbox"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e543f1e-bbcb-4276-9a09-f9181b2d86aa",
   "metadata": {},
   "source": [
    "### Define `tools.yaml`\n",
    "\n",
    "MCP Toolbox is configured using a YAML file that defines sources, tools, and toolsets.\n",
    "\n",
    "- Sources: Connections to your databases (e.g., Spanner, AlloyDB).\n",
    "- Tools: Specific actions the agent can take, like running a SQL query against a source.\n",
    "- Toolsets: Named groups of tools that can be loaded together by an agent.\n",
    "\n",
    "We use JSON format here because it's a valid subset of YAML and handles multi-line SQL strings cleanly.\n",
    "\n",
    "Notice the special syntax for the ${ALLOYDB_PASSWORD} environment variable below. This will be replaced by the value from Secret Manager mounted to the Cloud Run instance at runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05139b97-4933-45fe-98c7-4b78507434eb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Reference: https://googleapis.github.io/genai-toolbox/resources/sources/spanner/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/tools/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/tools/spanner-sql/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/sources/alloydb-pg/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/tools/postgres-sql/\n",
    "\n",
    "import os\n",
    "import json\n",
    "\n",
    "tools_config = {\n",
    "  \"sources\": {\n",
    "    \"spanner-finance-graph-source\": {\n",
    "      \"kind\": \"spanner\",\n",
    "      \"project\": f\"{project_id}\",\n",
    "      \"instance\": f\"{spanner_instance}\",\n",
    "      \"database\": f\"{spanner_database}\",\n",
    "      \"dialect\": \"googlesql\"\n",
    "    },\n",
    "    \"alloydb-finance-source\": {\n",
    "      \"kind\": \"alloydb-postgres\",\n",
    "      \"project\": f\"{project_id}\",\n",
    "      \"region\": f\"{region}\",\n",
    "      \"cluster\": f\"{alloydb_cluster}\",\n",
    "      \"instance\": f\"{alloydb_instance}\",\n",
    "      \"database\": f\"{alloydb_database}\",\n",
    "      \"user\": \"toolbox_user\",\n",
    "      \"password\": \"${ALLOYDB_PASSWORD}\",\n",
    "      \"ipType\": \"private\"\n",
    "    }\n",
    "  },\n",
    "  \"tools\": {\n",
    "    \"get_account_transfers\": {\n",
    "      \"kind\": \"spanner-sql\",\n",
    "      \"source\": \"spanner-finance-graph-source\",\n",
    "      \"description\": \"Use this tool to get information about the 10 most recent account transfers.\",\n",
    "      \"statement\": \"SELECT * FROM AccountTransferAccount LIMIT 10;\"\n",
    "    },\n",
    "    \"get_transactions\": {\n",
    "      \"kind\": \"postgres-sql\",\n",
    "      \"source\": \"alloydb-finance-source\",\n",
    "      \"description\": \"Use this tool to look up information about the 10 most recent credit card transactions.\",\n",
    "      \"statement\": \"SELECT * FROM transactions LIMIT 10;\"\n",
    "    }\n",
    "  },\n",
    "  \"toolsets\": {\n",
    "    \"finance-toolset\": [\n",
    "      \"get_account_transfers\",\n",
    "      \"get_transactions\"\n",
    "    ]\n",
    "  }\n",
    "}\n",
    "\n",
    "with open(\"tools.yaml\", \"w\") as file:\n",
    "    file.write(json.dumps(tools_config))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67efb5ce-bdd2-4855-b128-5806e4547cad",
   "metadata": {},
   "source": [
    "### Write `tools.yaml` to Secret Manager\n",
    "\n",
    "When deploying MCP Toolbox to Cloud Run, it is a best practice to write your tools.yaml file to a Secret Manager secret that can then be retrieved at runtime by Cloud Run. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c16dd0a7-d7e4-40e7-8463-e71a7d39f9ea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create the secret\n",
    "! gcloud secrets create tools --data-file=tools.yaml || gcloud secrets versions add tools --data-file=tools.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc5d2eec-c4cc-45cb-906e-2c65549b6bf1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Clean up the local file\n",
    "import os\n",
    "os.remove('tools.yaml')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34c4147c-44b3-4018-9991-71e9fc2c2cec",
   "metadata": {},
   "source": [
    "### Write AlloyDB Password to Secret Manager\n",
    "\n",
    "We can securely store the database password for the least-privilege `toolbox_user` account in Secret Manager to be accessed by Cloud Run as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e110de6-c586-4d54-9486-12aadcaf8eac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create the secret\n",
    "! echo -n {alloydb_toolbox_user_password} | gcloud secrets create alloydb-password --data-file=-\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9df899d1-004c-421a-973d-5f1bafdf8e12",
   "metadata": {},
   "source": [
    "### Deploy MCP Toolbox to Cloud Run\n",
    "\n",
    "Here, we deploy the MCP Toolbox as a secure Cloud Run service. Since this Cloud Run service interacts with sensitive databases, we want to be mindful of security best practices during deployment. \n",
    "\n",
    "Key security configurations:\n",
    "- `--network={vpc}` / `--subnet={vpc}`: Deploys the service within our private VPC.\n",
    "- `--service-account=toolbox-service-account`: Uses a dedicated service account with least-privilege permissions.\n",
    "- `--no-allow-unauthenticated`: Enforces IAM authentication for all invocations.\n",
    "- `--set-secrets`: Securely mounts our tools configuration and database password from Secret Manager.\n",
    "- `--ingress=internal`: Restricts network traffic to internal sources only.\n",
    "- `--telemetry-gcp`: Enables built-in OpenTelemetry for logging, tracing, and metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea5e992a-803e-45a7-8f31-5735625b3863",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Reference: https://cloud.google.com/run/docs/authenticating/public#gcloud\n",
    "\n",
    "# Define Toolbox Container Image\n",
    "image = f'{region}-docker.pkg.dev/database-toolbox/toolbox/toolbox:latest'\n",
    "\n",
    "# Deploy to Cloud Run\n",
    "! gcloud run deploy toolbox --no-user-output-enabled \\\n",
    "    --image={image} \\\n",
    "    --network={vpc} \\\n",
    "    --subnet={vpc} \\\n",
    "    --region={region} \\\n",
    "    --no-allow-unauthenticated \\\n",
    "    --set-secrets=\"/app/tools.yaml=tools:latest,ALLOYDB_PASSWORD=alloydb-password:latest\" \\\n",
    "    --args=\"--tools-file=/app/tools.yaml\",\"--address=0.0.0.0\",\"--port=8080\",\"--telemetry-gcp\" \\\n",
    "    --vpc-egress=all-traffic \\\n",
    "    --ingress=internal \\\n",
    "    --service-account=toolbox-service-account \\\n",
    "    --min=1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98ae510c-fbfd-41f9-a324-0e53649170c1",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Test Tools\n",
    "\n",
    "Tools are generally imported by an AI Agent, which chooses the correct tool for a given request and uses them to retrieve context and take other actions to achieve a stated objective. However, it is also useful for development and testing purposes to manually invoke the tools. In this step, we'll use the `toolbox_core` library to manually invoke the tools without an Agent framework to ensure they are working as expected. This pattern is also useful for running integration tests as part of a CI/CD pipeline when you make changes to tools. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf689b66-1555-45bf-91ca-8e716025e0de",
   "metadata": {},
   "source": [
    "### Get the Toolbox Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8471812d-4c25-48d2-96e2-c7f4a0026b38",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "toolbox_url = ! gcloud run services describe toolbox --region {region} --format 'value(metadata.annotations.\"run.googleapis.com/urls\")'\n",
    "toolbox_url = json.loads(toolbox_url[0])[0]\n",
    "print(toolbox_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2687057-7562-41d6-8c93-157eb2116149",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Get Auth Token\n",
    "\n",
    "We required authenticated invocations of the Cloud Run service, so we first need to grab an auth token to use with our ToolboxClient. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdb096bd-df6a-4bdb-9d65-48dcef2fb89b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import urllib\n",
    "\n",
    "import google.auth.transport.requests\n",
    "import google.oauth2.id_token\n",
    "\n",
    "\n",
    "def get_auth_token(endpoint):\n",
    "    # Cloud Run uses your service's hostname as the `audience` value\n",
    "    # audience = 'https://my-cloud-run-service.run.app/'\n",
    "    # For Cloud Run, `endpoint` is the URL (hostname + path) receiving the request\n",
    "    # endpoint = 'https://my-cloud-run-service.run.app/my/awesome/url'\n",
    "    \n",
    "    auth_req = google.auth.transport.requests.Request()\n",
    "    id_token = google.oauth2.id_token.fetch_id_token(auth_req, endpoint)\n",
    "\n",
    "    return id_token\n",
    "\n",
    "auth_token = get_auth_token(toolbox_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0857f2e-a7fe-49e2-9f16-ed9e1b8d20e1",
   "metadata": {},
   "source": [
    "### Execute Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dda4028-0cb6-41c6-a89f-93492df346c8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from toolbox_core import ToolboxClient, auth_methods\n",
    "\n",
    "auth_token_provider = auth_methods.aget_google_id_token # can also use sync method\n",
    "\n",
    "async with ToolboxClient(\n",
    "    toolbox_url,\n",
    "    client_headers={\"Authorization\": f\"Bearer {auth_token}\"},\n",
    ") as client:\n",
    "    tools = await client.load_toolset(\"finance-toolset\")\n",
    "    for t in tools:\n",
    "        print(f\"\\nRunning tool: {t._ToolboxTool__url}\")\n",
    "        result = await t()\n",
    "        json_result = json.loads(result)\n",
    "        print(\"Tool result:\\n\")\n",
    "        print(json.dumps(json_result, indent=2))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6f1f519-5d22-4f44-9324-78955dfa2bb8",
   "metadata": {},
   "source": [
    "## Update Tools\n",
    "\n",
    "So far, we have two simple database tools that run static SQL statements. This is useful for wiring up and testing the environment, but they won't be very useful for a real-world use case. To make our tools more dynamic, we'll start by adding parameters to our SQL queries, which allows the Agent to lookup data for specific customers based on their ID. \n",
    "\n",
    "We will use parameterized queries in our tools to prevent SQL injections. Query parameters can be used as substitutes for arbitrary expressions. Parameters cannot be used as substitutes for identifiers, column names, table names, or other parts of the query. Basic parameters types include string, integer, float, boolean types.\n",
    "\n",
    "Notice that AlloyDB and Spanner use a difference syntax for parameter substitution (i.e. `$1`, `$2`, etc. for AlloyDB and `@param_name1`, `@param_name2`, etc. for Spanner). Details on the specific syntax for each tool type can be found [in the docs](https://googleapis.github.io/genai-toolbox/resources/tools/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acdb0f6f-39c0-4ae0-8f2b-34a9bfc6ded2",
   "metadata": {},
   "source": [
    "### Update the `tools.yaml` File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e6dee9b-275b-45bd-b688-546ccc766d10",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Reference: https://googleapis.github.io/genai-toolbox/resources/sources/spanner/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/tools/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/tools/spanner-sql/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/sources/alloydb-pg/\n",
    "#            https://googleapis.github.io/genai-toolbox/resources/tools/postgres-sql/\n",
    "\n",
    "import os\n",
    "import json\n",
    "\n",
    "tools_config = {\n",
    "  \"sources\": {\n",
    "    \"spanner-finance-graph-source\": {\n",
    "      \"kind\": \"spanner\",\n",
    "      \"project\": f\"{project_id}\",\n",
    "      \"instance\": f\"{spanner_instance}\",\n",
    "      \"database\": f\"{spanner_database}\",\n",
    "      \"dialect\": \"googlesql\"\n",
    "    },\n",
    "    \"alloydb-finance-source\": {\n",
    "      \"kind\": \"alloydb-postgres\",\n",
    "      \"project\": f\"{project_id}\",\n",
    "      \"region\": f\"{region}\",\n",
    "      \"cluster\": f\"{alloydb_cluster}\",\n",
    "      \"instance\": f\"{alloydb_instance}\",\n",
    "      \"database\": f\"{alloydb_database}\",\n",
    "      \"user\": \"toolbox_user\",\n",
    "      \"password\": \"${ALLOYDB_PASSWORD}\",\n",
    "      \"ipType\": \"private\"\n",
    "    }\n",
    "  },\n",
    "  \"tools\": {\n",
    "    \"get_account_transfers\": {\n",
    "      \"kind\": \"spanner-sql\",\n",
    "      \"source\": \"spanner-finance-graph-source\",\n",
    "      \"description\": \"Use this tool to get information about the 10 most recent account transfers.\",\n",
    "      \"statement\": \"\"\"\n",
    "        SELECT p.id AS customer_id, *\n",
    "        FROM Person p\n",
    "          JOIN PersonOwnAccount poa ON poa.id = p.id\n",
    "          JOIN Account a ON a.id = poa.account_id\n",
    "          FULL JOIN AccountTransferAccount ata ON (ata.id = poa.account_id OR ata.to_id = poa.account_id)\n",
    "        WHERE p.id = @customer_id;\n",
    "        \"\"\",\n",
    "      \"parameters\": [\n",
    "        {\n",
    "          \"name\": \"customer_id\",\n",
    "          \"type\": \"integer\",\n",
    "          \"description\": \"Unique customer id number\"\n",
    "        }\n",
    "      ]\n",
    "    },\n",
    "    \"get_transactions\": {\n",
    "      \"kind\": \"postgres-sql\",\n",
    "      \"source\": \"alloydb-finance-source\",\n",
    "      \"description\": \"Use this tool to look up information about the 10 most recent credit card transactions.\",\n",
    "      \"statement\": \"\"\"\n",
    "        SELECT u.id AS customer_id, * \n",
    "        FROM users u\n",
    "        JOIN transactions t ON t.client_id = u.id\n",
    "        WHERE u.id = $1\n",
    "        ORDER BY date DESC\n",
    "        LIMIT 10\n",
    "        \"\"\",\n",
    "      \"parameters\": [\n",
    "        {\n",
    "          \"name\": \"customer_id\",\n",
    "          \"type\": \"integer\",\n",
    "          \"description\": \"Unique customer id number\"\n",
    "        }\n",
    "      ]\n",
    "    }\n",
    "  },\n",
    "  \"toolsets\": {\n",
    "    \"finance-toolset\": [\n",
    "      \"get_account_transfers\",\n",
    "      \"get_transactions\"\n",
    "    ]\n",
    "  }\n",
    "}\n",
    "\n",
    "with open(\"tools.yaml\", \"w\") as file:\n",
    "    file.write(json.dumps(tools_config))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09a2d9cf-8d4a-4beb-b1b7-f8b1dbc5623d",
   "metadata": {},
   "source": [
    "### Write Updated `tools.yaml` to Secret Manager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52fa42e6-980e-435d-959c-4834ead46faa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create the secret\n",
    "! gcloud secrets versions add tools --data-file=tools.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37610a56-b240-41e8-8f75-1f897a82696d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Clean up the local file\n",
    "import os\n",
    "os.remove('tools.yaml')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "895f51c6-c70c-4840-b5fa-2fb22a68c122",
   "metadata": {},
   "source": [
    "### Update Toolbox with New `tools.yaml` File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cf8fee5-2087-4072-98ab-b92a5f6892d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Reference: https://cloud.google.com/sdk/gcloud/reference/run/services/update\n",
    "\n",
    "! gcloud run services update toolbox --no-user-output-enabled \\\n",
    "    --update-secrets=\"/app/tools.yaml=tools:latest,ALLOYDB_PASSWORD=alloydb-password:latest\" \\\n",
    "    --region={region}\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d15e06-ead0-4942-a165-defa825c8a42",
   "metadata": {},
   "source": [
    "### Execute Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48ce9cdd-1804-401c-b211-1166f1d6e760",
   "metadata": {},
   "source": [
    "#### Run Without `customer_id`\n",
    "\n",
    "We're first run the updated tools without providing a customer ID. Notice that this fails because this is a required parameter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7eb42e03-3404-4054-8416-0728984ecaf9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from toolbox_core import ToolboxClient, auth_methods\n",
    "\n",
    "# Get endpoint\n",
    "toolbox_url = ! gcloud run services describe toolbox --region {region} --format 'value(metadata.annotations.\"run.googleapis.com/urls\")'\n",
    "toolbox_url = json.loads(toolbox_url[0])[0]\n",
    "print(f\"Toolbox Cloud Run endpoint: {toolbox_url}\")\n",
    "\n",
    "# Refresh auth_token\n",
    "auth_token = get_auth_token(toolbox_url)\n",
    "auth_token_provider = auth_methods.aget_google_id_token # can also use sync method\n",
    "\n",
    "# Run tools \n",
    "async with ToolboxClient(\n",
    "    toolbox_url,\n",
    "    client_headers={\"Authorization\": f\"Bearer {auth_token}\"},\n",
    ") as client:\n",
    "    tools = await client.load_toolset(\"finance-toolset\")\n",
    "    for t in tools:\n",
    "        print(f\"\\nRunning tool: {t._ToolboxTool__url}\")\n",
    "        result = await t()\n",
    "        json_result = json.loads(result)\n",
    "        print(\"Tool result:\\n\")\n",
    "        print(json.dumps(json_result, indent=2))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b09a259-c30a-45ae-aa9f-7589c588d2f1",
   "metadata": {},
   "source": [
    "#### Run With `customer_id`\n",
    "\n",
    "We'll update our code to provide the required parameter this time. When Agents execute tools, they will be aware of the required parameters and can populate them in their tool calls based on the session context. Try setting different values (between 1 and 200) for customer_id and observe that the tools now return  account and transaction information about specific customers.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50745d7b-e6b5-49c2-a938-ed06c7aacd68",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from toolbox_core import ToolboxClient, auth_methods\n",
    "\n",
    "# Set customer_id \n",
    "customer_id=11\n",
    "\n",
    "# Refresh auth_token\n",
    "auth_token = get_auth_token(toolbox_url)\n",
    "auth_token_provider = auth_methods.aget_google_id_token # can also use sync method\n",
    "\n",
    "# Run tools \n",
    "async with ToolboxClient(\n",
    "    toolbox_url,\n",
    "    client_headers={\"Authorization\": f\"Bearer {auth_token}\"},\n",
    ") as client:\n",
    "    tools = await client.load_toolset(\"finance-toolset\")\n",
    "    for t in tools:\n",
    "        print(f\"\\nRunning tool: {t._ToolboxTool__url}\")\n",
    "        result = await t(customer_id) # Call the tool with the customer_id parameter here\n",
    "        json_result = json.loads(result)\n",
    "        print(\"Tool result:\\n\")\n",
    "        print(json.dumps(json_result, indent=2))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "175c35ba",
   "metadata": {},
   "source": [
    "Congratulations, you have completed Module 2! Proceed to [`3_build_adk_agent.ipynb`](./3_build_adk_agent.ipynb) to build your first ADK Agent and connect it to your database tools."
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m130",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m130"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
